Skip to content

feat(rocm): MoE combine/gate ops — SharedExpertGate, MoeCombine, MoeCombineGate (issue #41) - #509

Draft
VikashLoomba wants to merge 3 commits into
mudler:mainfrom
VikashLoomba:row/ROCM-MOE-CHAIN
Draft

feat(rocm): MoE combine/gate ops — SharedExpertGate, MoeCombine, MoeCombineGate (issue #41)#509
VikashLoomba wants to merge 3 commits into
mudler:mainfrom
VikashLoomba:row/ROCM-MOE-CHAIN

Conversation

@VikashLoomba

@VikashLoomba VikashLoomba commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Row

BACKEND-ROCM — the next links in the generic MoE path, after the router/silu-mul (#348). Issue #41. Claim CLAIM-ROCM-GDN-KERNELS continues.

What changed

NEW src/vt/rocm/rocm_moe_chain.hip with three ops, hand-translated from cuda_moe.cu (MoeCombineKernel :473, MoeCombineGateKernel :555) + the SharedExpertGate CPU oracle (cpu_ops.cpp:2387):

  • kSharedExpertGate — out[t,c] = sigmoid(gl[t]) * sd[t,c] (bf16 out, f32/bf16 sd)
  • kMoeCombine — weighted top-k expert sum + optional shared term
  • kMoeCombineGate — combine with the shared-expert sigmoid gate folded in, the shared term rounded through bf16 exactly as the donor

All grid-stride, f32 math, bf16/f32 dtype arms via boundary conversions. New cross-device case gates all three (MoeCombineGate's oracle is the host-computed composite — no CPU op registration exists for it).

Evidence (4× gfx1100, ROCm 7.14, Release)

  • MoE combine/gate case: 9/9 assertions, runs not skips
  • ctest -R 'rocm|cross_device': 4/4
  • full ctest: pre-existing failure set shrinks 7 → 5 — test_bench and test_capi now PASS (they failed at op 77 / the router dtype before the chain). test_loaded_engine_dense now fails only on the async-scheduling assertion (runner_supports_async()=false on ROCm), a lane capability gap, not a kernel throw.
  • preflight --staged + trailers green

Speed claims

  • This PR makes NO speed claim.

Honest gaps

  • Named remaining blocker: the grouped quant expert GEMM (kMatmulBTQuantGrouped, the DeepSeek-V4 keep-quant family) — the heavy lift for Qwen3.5-27B-class GDN-MoE models, a proper kernel project of its own.
  • The router's grouped/bias/noaux_tc forms still throw by design (same guard as feat(rocm): gfx1201 hipBLAS ops + Gemma-4-26B-A4B MoE (BF16/FP8) #140); the covered path is ungrouped softmax.

@VikashLoomba VikashLoomba mentioned this pull request Aug 12, 2026
Draft
1 task
@joral joral mentioned this pull request Aug 13, 2026
Merged
@localai-bot localai-bot mentioned this pull request Aug 13, 2026
Draft
1 task
@localai-bot

localai-bot commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Reviewed as part of a sweep over the open external PRs. The translations are faithful — I checked all three donor anchors at your base SHA and they resolve exactly (cuda_moe.cu:473 is MoeCombineKernel, :555 is MoeCombineGateKernel, cpu_ops.cpp:2387 is SharedExpertGateKernel), and the kernel bodies match the donors call for call, including the __bfloat162float(__float2bfloat16(sv)) double-round. Registration goes through the existing registrar with no parallel path.

Two things to fix. Both are the same shape as findings on #506 and #523 — the calculation ported cleanly, the guards around it did not — so it is probably worth reading the three together.

1. The donor's dtype refusals were dropped, and f16 is reachable.

MoeCombineKernelCuda (cuda_moe.cu:520-524) and MoeCombineGateKernelCuda (:597-604) each open with VT_CHECKs refusing anything but f32/bf16, with a named message. The ROCm entry points have none, and dispatch is a bare if (expert_out.dtype == DType::kBF16) ... else <float>.

f16 gets through: vt::MoeCombine gates on IsFloat (src/vt/ops.cpp:21 — kF32 || kF16 || kBF16) for both expert_out and shared, and Tensor::Ptr<T>() (include/vt/tensor.h:66-68) is an unchecked static_cast with no dtype assertion. So an f16 expert_out [T,K,H] passes the seam, falls into the else branch, and the kernel reads 4 bytes per element out of a 2-bytes-per-element allocation — T*K*H*2 bytes past the end, garbage out, no error anywhere. AGENTS.md is explicit that an unimplemented arm is refused with a message naming the missing piece.

2. The new test exercises the dtype arm the model never runs.

On the live path expert_out is bf16 (qwen3_5.cpp:5463, DBuf ddown(d, DType::kBF16, {P, H})), dout is bf16, and the unfused arm passes a bf16 shared (:5326). The new case in test_backend_cross_device.cpp only ever constructs DType::kF32 for expert_out, shared, sd and out, so MoeCombine and MoeCombineGate are tested exclusively as <float,float,float>. The comment says "f32 and bf16 arms"; only SharedExpertGate actually exercises a bf16 store.

Consequence: an inverted predicate in by_out/by_shared/by_sd, or a wrong bf16 Ld/St boundary, is a 2x out-of-bounds read on every MoE layer of every token — and the case still reports 9/9 green, because those branches are never entered.

Smaller ones: CMakeLists.txt:1396-1397 has src/vt/rocm/rocm_moe_chain.hip twice with mangled indentation (inert, but it is exactly the shape that merges badly against another additive ROCm PR touching the same list); docs/FEATURES.md:232,320 still say "44 registered ops" where the count is now 47; and the MoeCombine f32 arm is gated at NMSE 5e-4 where the donor asserts bit-exactness by design (cuda_moe.cu:465-468 — single store-rounding, same as the CPU reference) and the ROCm kernel is thread-per-element with no cross-lane reduction, so CHECK(got == ref) is achievable and is the stronger gate. Your SharedExpertGate check already does exactly that.

Credit where due: recomputing the MoeCombineGate oracle from scratch on host rather than routing it through a shared helper is the right call, and your stated reason for it (no CPU registration for kMoeCombineGate) checks out — cpu_ops.cpp registers kMoeCombine only.

No AMD hardware here, so your 9/9 and the ctest results could not be reproduced and I am not disputing them; both findings are static, read from the seam and the model path.

VikashLoomba added a commit to VikashLoomba/vllm.cpp that referenced this pull request Aug 14, 2026
@VikashLoomba
fix(ROCM): MoE combine/gate dtype refusals + production-dtype test ar…
dfd1213 
…ms -- the mudler#509 review rework

Review sweep findings (localai-bot, 2026-08-13), all accepted:

1. The donor's dtype refusals were dropped: cuda_moe.cu:520-524/:597-604 open
   with VT_CHECKs refusing non-f32/bf16; without them an f16 expert_out passes
   the seam's IsFloat gate and Tensor::Ptr<T>()'s unchecked cast reads 4 bytes
   per element from a 2-byte allocation. Refusals added to all three ROCm
   entry points (SharedExpertGate included: its 4-arm dispatch has the same
   f16 hazard on sd).
2. The case now exercises the PRODUCTION dtype mix, not only f32: the model
   path runs expert_out bf16 (qwen3_5.cpp DBuf ddown), shared bf16, out bf16.
   The new bf16 arm is asserted BIT-EXACT against the CPU reference (both
   sides thread-per-element, same sequential K order, single store rounding,
   -ffp-contract=off), and the f32 MoeCombine arm is tightened from NMSE to
   bit-exact per the donor's design comment (cuda_moe.cu:465-468). Writing the
   bf16 arm caught a construction bug in the first version of it (an f32-typed
   tensor over a bf16 buffer -- exactly the OOB class the review predicted the
   missing arm hid); fixed and re-verified against an independent host
   composite (0/320) plus a raw-hipMalloc scratch replica of both backends.
3. CMakeLists.txt: the mangled duplicate rocm_moe_chain.hip line removed.
4. docs/FEATURES.md op count 44 -> 47 (counted: the registration sites).

Gates (gfx1100, flock): test_backend_cross_device 20/20 (bit-exact bf16 +
f32 arms).

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: pi:kimi-k3 [pi]
@VikashLoomba

VikashLoomba commented Aug 14, 2026

Copy link
Copy Markdown
Contributor Author

Both findings accepted and reworked (commits dfd1213 + ebb4c6b, rebased onto current main):

  1. dtype refusals restored: all three ROCm entry points now open with the donor's VT_CHECKs (f32/bf16 only, named messages) — the f16-through-IsFloat → unchecked-Ptr<T> OOB read is refused loudly. SharedExpertGate got the same treatment (its 4-arm dispatch had the same hazard on sd).
  2. Production dtype arms: the case now exercises the mix the model actually runs (expert_out bf16 / shared bf16 / out bf16, per qwen3_5.cpp) in addition to f32 — and both MoeCombine arms are asserted bit-exact, not NMSE (thread-per-element, no cross-lane reduction, -ffp-contract=off; the donor's design comment says the same). Writing the bf16 arm immediately caught a construction bug in my first version of it (an f32-typed tensor over a bf16 buffer — the exact class you predicted the missing arm hid), which I root-caused against an independent host composite before believing either side.

Also: the mangled duplicate rocm_moe_chain.hip line in CMakeLists.txt is removed; FEATURES.md op count corrected 44 → 47 (counted the registration sites); and the rebase onto main surfaced main's new MoeCombineFn(..., float routed_scale) — plumbed through the ROCm kernel in the CPU reference's order (scale the routed sum, then add shared), with the bf16 arm now running at scale 0.7 so the multiply is exercised.

Gates: test_backend_cross_device 20/20 on gfx1100.

VikashLoomba added a commit to VikashLoomba/vllm.cpp that referenced this pull request Aug 14, 2026
@VikashLoomba
fix(ROCM): MoE combine/gate dtype refusals + production-dtype test ar…
5e68013 
…ms -- the mudler#509 review rework

Review sweep findings (localai-bot, 2026-08-13), all accepted:

1. The donor's dtype refusals were dropped: cuda_moe.cu:520-524/:597-604 open
   with VT_CHECKs refusing non-f32/bf16; without them an f16 expert_out passes
   the seam's IsFloat gate and Tensor::Ptr<T>()'s unchecked cast reads 4 bytes
   per element from a 2-byte allocation. Refusals added to all three ROCm
   entry points (SharedExpertGate included: its 4-arm dispatch has the same
   f16 hazard on sd).
2. The case now exercises the PRODUCTION dtype mix, not only f32: the model
   path runs expert_out bf16 (qwen3_5.cpp DBuf ddown), shared bf16, out bf16.
   The new bf16 arm is asserted BIT-EXACT against the CPU reference (both
   sides thread-per-element, same sequential K order, single store rounding,
   -ffp-contract=off), and the f32 MoeCombine arm is tightened from NMSE to
   bit-exact per the donor's design comment (cuda_moe.cu:465-468). Writing the
   bf16 arm caught a construction bug in the first version of it (an f32-typed
   tensor over a bf16 buffer -- exactly the OOB class the review predicted the
   missing arm hid); fixed and re-verified against an independent host
   composite (0/320) plus a raw-hipMalloc scratch replica of both backends.
3. CMakeLists.txt: the mangled duplicate rocm_moe_chain.hip line removed.
4. docs/FEATURES.md op count 44 -> 47 (counted: the registration sites).

Gates (gfx1100, flock): test_backend_cross_device 20/20 (bit-exact bf16 +
f32 arms).

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: pi:kimi-k3 [pi]
@joral joral mentioned this pull request Aug 14, 2026
Open
VikashLoomba added a commit to VikashLoomba/vllm.cpp that referenced this pull request Aug 14, 2026
@VikashLoomba
fix(ROCM): MoE combine/gate dtype refusals + production-dtype test ar…
72bb347 
…ms -- the mudler#509 review rework

Review sweep findings (localai-bot, 2026-08-13), all accepted:

1. The donor's dtype refusals were dropped: cuda_moe.cu:520-524/:597-604 open
   with VT_CHECKs refusing non-f32/bf16; without them an f16 expert_out passes
   the seam's IsFloat gate and Tensor::Ptr<T>()'s unchecked cast reads 4 bytes
   per element from a 2-byte allocation. Refusals added to all three ROCm
   entry points (SharedExpertGate included: its 4-arm dispatch has the same
   f16 hazard on sd).
2. The case now exercises the PRODUCTION dtype mix, not only f32: the model
   path runs expert_out bf16 (qwen3_5.cpp DBuf ddown), shared bf16, out bf16.
   The new bf16 arm is asserted BIT-EXACT against the CPU reference (both
   sides thread-per-element, same sequential K order, single store rounding,
   -ffp-contract=off), and the f32 MoeCombine arm is tightened from NMSE to
   bit-exact per the donor's design comment (cuda_moe.cu:465-468). Writing the
   bf16 arm caught a construction bug in the first version of it (an f32-typed
   tensor over a bf16 buffer -- exactly the OOB class the review predicted the
   missing arm hid); fixed and re-verified against an independent host
   composite (0/320) plus a raw-hipMalloc scratch replica of both backends.
3. CMakeLists.txt: the mangled duplicate rocm_moe_chain.hip line removed.
4. docs/FEATURES.md op count 44 -> 47 (counted: the registration sites).

Gates (gfx1100, flock): test_backend_cross_device 20/20 (bit-exact bf16 +
f32 arms).

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: pi:kimi-k3 [pi]
@VikashLoomba
feat(rocm): MoE combine/gate ops (SharedExpertGate, MoeCombine, MoeCo…
25b9485 
…mbineGate) (mudler#41)

The next links in the generic MoE path after the router/silu-mul. Hand-
translated from cuda_moe.cu (MoeCombineKernel :473, MoeCombineGateKernel :555)
and the SharedExpertGate CPU oracle (cpu_ops.cpp:2387). Grid-stride, f32 math,
bf16/f32 dtype arms via boundary conversions; the combine-gate folds the
shared-expert sigmoid gate rounded through bf16 exactly as the donor.

Evidence (4x gfx1100, ROCm 7.14, Release):
- new MoE combine/gate cross-device case: 9/9 assertions (MoeCombineGate's
  oracle is the host-computed composite — no CPU op registration exists)
- ctest -R 'rocm|cross_device': 4/4
- full ctest: pre-existing failure set shrinks 7 -> 5; test_bench and
  test_capi now PASS (they failed at op 77 / the router dtype before the
  chain). test_loaded_engine_dense now fails only on the async-scheduling
  assertion (a lane capability gap, not a kernel throw).
- Named remaining blocker: the grouped quant expert GEMM
  (kMatmulBTQuantGrouped), the DeepSeek-V4 keep-quant family.

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: pi:kimi-k3 [pi]
@VikashLoomba
fix(ROCM): MoE combine/gate dtype refusals + production-dtype test ar…
202a37c 
…ms -- the mudler#509 review rework

Review sweep findings (localai-bot, 2026-08-13), all accepted:

1. The donor's dtype refusals were dropped: cuda_moe.cu:520-524/:597-604 open
   with VT_CHECKs refusing non-f32/bf16; without them an f16 expert_out passes
   the seam's IsFloat gate and Tensor::Ptr<T>()'s unchecked cast reads 4 bytes
   per element from a 2-byte allocation. Refusals added to all three ROCm
   entry points (SharedExpertGate included: its 4-arm dispatch has the same
   f16 hazard on sd).
2. The case now exercises the PRODUCTION dtype mix, not only f32: the model
   path runs expert_out bf16 (qwen3_5.cpp DBuf ddown), shared bf16, out bf16.
   The new bf16 arm is asserted BIT-EXACT against the CPU reference (both
   sides thread-per-element, same sequential K order, single store rounding,
   -ffp-contract=off), and the f32 MoeCombine arm is tightened from NMSE to
   bit-exact per the donor's design comment (cuda_moe.cu:465-468). Writing the
   bf16 arm caught a construction bug in the first version of it (an f32-typed
   tensor over a bf16 buffer -- exactly the OOB class the review predicted the
   missing arm hid); fixed and re-verified against an independent host
   composite (0/320) plus a raw-hipMalloc scratch replica of both backends.
3. CMakeLists.txt: the mangled duplicate rocm_moe_chain.hip line removed.
4. docs/FEATURES.md op count 44 -> 47 (counted: the registration sites).

Gates (gfx1100, flock): test_backend_cross_device 20/20 (bit-exact bf16 +
f32 arms).

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: pi:kimi-k3 [pi]
@VikashLoomba
fix(ROCM): MoeCombine routed_scale (main's signature change at mudler…
9aefaec 
…#684) plumbed through the ROCm arm

The rebase onto current main brought main's new MoeCombineFn signature
(routed_scale, default 1.0f, scaling the ROUTED sum before the shared term --
upstream apply_routed_scale_to_output). The ROCm kernel applies it in the same
f32 accumulator in the same order (one standalone multiply on the finished
sum, bit-identical to the CPU reference under -ffp-contract=off); the forward
declaration in rocm_ops.hip is updated to match. The bf16 test arm now runs
at scale 0.7 so the multiply is exercised, not just the 1.0 passthrough.

Gates (gfx1100, flock): test_backend_cross_device 20/20 (bit-exact arms
unchanged at 1.0, bit-exact at 0.7).

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: pi:kimi-k3 [pi]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@VikashLoomba @localai-bot